#!/usr/bin/env python3
"""
快速演示脚本 - 直接运行圆柱形滚筒优化
"""

import os
import sys
import numpy as np
import torch
import matplotlib.pyplot as plt
import json

# 添加路径
sys.path.append('inverse_barrier')

def quick_demo():
    """快速演示函数"""
    print("Cylinder drum demo")
    print("="*40)
    
    # 创建输出目录
    output_dir = "inverse_barrier/data/outputs_cylinder"
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        # 1. 测试圆柱几何体生成
        print("Generating cylinder geometry...")
        from utils_cylinder import create_cylinder_drum
        
        cylinder_particles = create_cylinder_drum(
            inner_radius=0.05,    # 50mm
            outer_radius=0.055,   # 55mm  
            length_x=0.03,        # 30mm
            spacing=0.004,        # 4mm
            include_end_caps=True
        )
        
        print(f"   Cylinder particles: {len(cylinder_particles)}")
        
        # 2. 可视化圆柱几何体
        print("Creating visualizations...")
        particles_np = cylinder_particles.detach().numpy()
        
        fig = plt.figure(figsize=(12, 4))
        
        # 3D视图
        ax1 = fig.add_subplot(1, 3, 1, projection='3d')
        ax1.scatter(particles_np[:, 0], particles_np[:, 1], particles_np[:, 2], 
                   s=0.5, alpha=0.6, c='blue')
        ax1.set_xlabel('X')
        ax1.set_ylabel('Y')
        ax1.set_zlabel('Z')
        ax1.set_title('Cylinder 3D View')
        
        # Y-Z截面
        ax2 = fig.add_subplot(1, 3, 2)
        ax2.scatter(particles_np[:, 1], particles_np[:, 2], s=1, alpha=0.6, c='blue')
        
        # 绘制内外径参考圆
        theta = np.linspace(0, 2*np.pi, 100)
        ax2.plot(0.05 * np.cos(theta), 0.05 * np.sin(theta), 'r--', linewidth=1, label='Inner')
        ax2.plot(0.055 * np.cos(theta), 0.055 * np.sin(theta), 'g--', linewidth=1, label='Outer')
        
        ax2.set_xlabel('Y (m)')
        ax2.set_ylabel('Z (m)')
        ax2.set_title('Cross Section')
        ax2.set_aspect('equal')
        ax2.legend()
        ax2.grid(True)
        
        # X方向分布
        ax3 = fig.add_subplot(1, 3, 3)
        ax3.hist(particles_np[:, 0], bins=20, alpha=0.7, color='blue')
        ax3.axvline(0, color='r', linestyle='--', label='Left')
        ax3.axvline(0.03, color='r', linestyle='--', label='Right')
        ax3.set_xlabel('X (m)')
        ax3.set_ylabel('Count')
        ax3.set_title('X Distribution')
        ax3.legend()
        ax3.grid(True)
        
        plt.tight_layout()
        plt.savefig(f'{output_dir}/cylinder_geometry_demo.png', dpi=150, bbox_inches='tight')
        plt.close()
        
        print(f"   Saved: {output_dir}/cylinder_geometry_demo.png")
        
        # 3. 简单的优化演示
        print("Running optimization demo...")
        
        # 生成测试颗粒
        n_particles = 200
        inner_radius = 0.05
        
        # 在圆柱内部随机生成颗粒
        test_particles = []
        for _ in range(n_particles):
            r = np.random.uniform(0, inner_radius * 0.8)
            theta = np.random.uniform(0, 2*np.pi)
            x = np.random.uniform(0, 0.03)
            y = r * np.cos(theta)
            z = r * np.sin(theta)
            test_particles.append([x, y, z])
        
        test_particles = np.array(test_particles)
        
        # 测试不同旋转角度的效果
        angles = np.linspace(0, np.pi, 8)
        results = []
        
        for angle in angles:
            # 简单的旋转变换
            cos_a, sin_a = np.cos(angle), np.sin(angle)
            rotated_y = test_particles[:, 1] * cos_a - test_particles[:, 2] * sin_a
            rotated_z = test_particles[:, 1] * sin_a + test_particles[:, 2] * cos_a
            
            # 计算一个简单的目标函数（例如：Y方向的分散度）
            dispersion = np.std(rotated_y)
            results.append(dispersion)
        
        # 绘制优化结果
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
        
        # 优化曲线
        ax1.plot(angles * 180 / np.pi, results, 'bo-', linewidth=2, markersize=6)
        ax1.set_xlabel('Rotation Angle (deg)')
        ax1.set_ylabel('Y Dispersion')
        ax1.set_title('Optimization Results')
        ax1.grid(True)
        
        # 标记最优点
        best_idx = np.argmin(results)
        best_angle = angles[best_idx]
        ax1.plot(best_angle * 180 / np.pi, results[best_idx], 'ro', markersize=10, 
                label=f'Optimal: {best_angle*180/np.pi:.1f}°')
        ax1.legend()
        
        # 最优角度下的颗粒分布
        cos_best, sin_best = np.cos(best_angle), np.sin(best_angle)
        best_rotated_y = test_particles[:, 1] * cos_best - test_particles[:, 2] * sin_best
        best_rotated_z = test_particles[:, 1] * sin_best + test_particles[:, 2] * cos_best
        
        ax2.scatter(best_rotated_y, best_rotated_z, s=20, alpha=0.6, c='red', label='Rotated')
        ax2.scatter(test_particles[:, 1], test_particles[:, 2], s=20, alpha=0.6, c='blue', label='Original')
        
        # 绘制圆柱边界
        theta = np.linspace(0, 2*np.pi, 100)
        ax2.plot(inner_radius * np.cos(theta), inner_radius * np.sin(theta), 'k-', linewidth=2, label='Wall')
        
        ax2.set_xlabel('Y (m)')
        ax2.set_ylabel('Z (m)')
        ax2.set_title(f'Optimal Angle: {best_angle*180/np.pi:.1f}°')
        ax2.set_aspect('equal')
        ax2.legend()
        ax2.grid(True)
        
        plt.tight_layout()
        plt.savefig(f'{output_dir}/optimization_demo.png', dpi=150, bbox_inches='tight')
        plt.close()
        
        print(f"   Saved: {output_dir}/optimization_demo.png")
        
        # 4. 保存结果摘要
        summary = {
            'specs': {
                'inner_dia_mm': 100,
                'outer_dia_mm': 110,
                'length_mm': 30,
                'wall_thickness_mm': 5
            },
            'results': {
                'cylinder_particles': len(cylinder_particles),
                'test_particles': n_particles,
                'optimal_angle_deg': float(best_angle * 180 / np.pi),
                'optimal_angle_rad': float(best_angle),
                'min_dispersion': float(results[best_idx])
            }
        }
        
        with open(f'{output_dir}/summary.json', 'w') as f:
            json.dump(summary, f, indent=2)
        
        print(f"   Saved: {output_dir}/summary.json")
        
        print("\nResults:")
        print(f"Cylinder: {summary['specs']['inner_dia_mm']}mm ID, {summary['specs']['outer_dia_mm']}mm OD, {summary['specs']['length_mm']}mm L")
        print(f"Particles: {summary['results']['cylinder_particles']} wall, {summary['results']['test_particles']} test")
        print(f"Optimal angle: {summary['results']['optimal_angle_deg']:.1f}°")
        print(f"Min dispersion: {summary['results']['min_dispersion']:.4f}")
        print(f"Output: {output_dir}")
        
        return True
        
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = quick_demo()
    if success:
        print("Demo completed successfully.")
    else:
        print("Demo failed.") 